package com.shujia.spark.sql

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object Demo06RDDToDF {
  case class Student(id: String, name: String, age: Int, gender: String, clz: String)

  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      .master("local")
      .config("spark.sql.shuffle.partitions", "2")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    // RDD 与 DF之间的相互转换
    val stuDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("id String,name String,age Int,gender String,clazz String")
      .load("spark/data/stu/students.txt")

    // DF to RDD
    val stuRDD: RDD[Row] = stuDF.rdd
    // 通过getAs[泛型]("字段名")方法可以从Row对象中提取某一列的数据
    // 过滤性别为女的理科五班的同学
    stuRDD
      .filter(row => {
        var flag: Boolean = false
        val gender: String = row.getAs[String]("gender")
        val clazz: String = row.getAs[String]("clazz")
        if ("女".equals(gender) && "理科五班".equals(clazz)) {
          flag = true
        }
        flag
      }).map(row => {
        val id: String = row.getAs[String]("id")
        val name: String = row.getAs[String]("name")
        val age: Int = row.getAs[Int]("age")
        s"$id,$name,$age"
      }).foreach(println)

    val dfToRDD02: RDD[(String, String, Int, String, String)] = stuDF
      .as[(String, String, Int, String, String)]
      .rdd

    dfToRDD02.map {
      case (id: String, name: String, age: Int, gender: String, clazz: String) =>
        s"$id,$name,$age"
    }.foreach(println)

    // Rdd to DF
    val stuLineRDD: RDD[String] = spark
      .sparkContext.textFile("spark/data/stu/students.txt")

    val rddToDF01: DataFrame = stuLineRDD.map(line => {
      val splits: Array[String] = line.split(",")
      (splits(0), splits(1), splits(2).toInt, splits(3), splits(4))
    }).toDF("id", "name", "age", "gender", "clazz")

    rddToDF01.show()

    // 通过样例类对象进行DF的转换
    val rddToDF02: DataFrame = stuLineRDD.map(line => {
      val splits: Array[String] = line.split(",")
      Student(splits(0), splits(1), splits(2).toInt, splits(3), splits(4))
    }).toDF()

    rddToDF02.show()
  }

}
